import numpy as np
import os
import output_function as of
import pandas as pd
from operatorset import compute_stream_diriclet

def vel_to_omega(u, v, dx, dy):  
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega

#扩散算符
def diff_operator(myfun,nu,dx,dy):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x+diff_y



# 使用示例
base_path = './data/RB_cavity/'  # 根据需要修改为您的保存路径
out_iter =20000  # 假设我们要加载第300时间步的数据

# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')

n_density=n

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
Ra=Ra_star = config['Ra_star']
prandtl=Pr = config['Pr']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 打印一些变量以确认它们已被正确加载和赋值
print(f"dx: {dx}")
print(f"Ra_star: {Ra_star}")
print(f"dt: {dt}")
# 确认是否成功加载了正确的时间步长
if out_iter_loaded == out_iter:
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    print(omega.shape)
else:
    print("Failed to load the correct time step or file does not exist.")

# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
print(xx)
# error=vel_to_omega(u,v,dx,dy)-omega[1:-1,1:-1]
# print(error)

psi, end_signal = compute_stream_diriclet(omega[1:-1, 1:-1],Nx,Ny, dx, dy, tol=1e-7, psi_init=psi)

omega_numerical=-diff_operator(psi,1,dx,dy)
error=omega[1:-1,1:-1]-omega_numerical
errm=np.max(np.abs(error))
print(errm)


# data = pd.DataFrame(error)
# writer = pd.ExcelWriter(f'{base_path}omega_error_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# data = pd.DataFrame(omega)
# writer = pd.ExcelWriter(f'{base_path}omega_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# 可视化结果（可选）
import matplotlib.pyplot as plt
# fig, axs = plt.subplots(1, 3, figsize=(12, 6))
# cax1 = axs[0].imshow(omega_numerical, extent=[0, Lx, 0, Ly], cmap='viridis', origin='lower')
# axs[0].set_title('Numerical Solution')
# fig.colorbar(cax1, ax=axs[0])

# cax2 = axs[1].imshow(omega, extent=[0, Lx, 0, Ly], cmap='viridis', origin='lower')
# axs[1].set_title('Analytic Solution')
# fig.colorbar(cax2, ax=axs[1])

# cax3 = axs[2].imshow(error, extent=[0, Lx, 0, Ly], cmap='viridis', origin='lower')
# axs[2].set_title('Analytic Solution')
# fig.colorbar(cax3, ax=axs[2])

# plt.show()

# 绘制数值解、解析解和误差
plt.figure(figsize=(15, 5))

# 数值解
plt.subplot(1, 3, 1)
plt.title("Numerical Solution")
plt.imshow(omega_numerical, extent=[0, Lx, 0, Ly], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 精确解
plt.subplot(1, 3, 2)
plt.title("Analytical Solution")
plt.imshow(omega, extent=[0, Lx, 0, Ly], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 误差
plt.subplot(1, 3, 3)
plt.title("Error between Numerical and Analytical Solutions")
plt.imshow(error, extent=[0, Lx, 0, Ly], origin='lower', cmap='coolwarm')
plt.colorbar(label='Error')

plt.tight_layout()
plt.show()



# data = pd.DataFrame(u)
# writer = pd.ExcelWriter(f'{base_path}u_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# data = pd.DataFrame(v)
# writer = pd.ExcelWriter(f'{base_path}v_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()
    
# omega_drive_term=-Ra*prandtl*(n_density[1:-1,2:]-n_density[1:-1,:-2])/2/dx
# #print(omega_drive_term)
# data = pd.DataFrame(omega_drive_term)
# writer = pd.ExcelWriter(f'{base_path}omega_drive_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# def diff_operator(myfun,nu):
#     diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
#     diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
#     return diff_x+diff_y

# omega_diff=diff_operator(omega,Pr)
# data = pd.DataFrame(omega_diff)
# writer = pd.ExcelWriter(f'{base_path}omega_diff_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# 使用加载的数据进行分析、绘图或其他处理
# base_folder=f'./RB_cavity_plot_Ra_{Ra_star}_Lx_{Lx}'
# of.diag_n_pcolor_filepath(n, n_init0_matrix, u, v, xx, yy, out_iter, Ra_star, dt,base_folder)
# of.plot_and_save_n_density_x_section(x,y, n, n_init0_matrix, Ny//2, out_iter,title_base="delta n_density Distribution",save_folder_base=base_folder)
# of.plot_and_save_v_x_section(x,y, v, Ny//2, out_iter,title_base="V Distribution along X-axis",save_folder_base=base_folder)